package inference.traffic;

import java.io.File;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Random;
import java.util.Set;

import rddl.EvalException;
import rddl.RDDL;
import rddl.RDDL.CPF_DEF;
import rddl.RDDL.DOMAIN;
import rddl.RDDL.EXPR;
import rddl.RDDL.INSTANCE;
import rddl.RDDL.LCONST;
import rddl.RDDL.LVAR;
import rddl.RDDL.NONFLUENTS;
import rddl.RDDL.PVAR_INST_DEF;
import rddl.RDDL.PVAR_NAME;
import rddl.RDDL.TYPE_NAME;
import rddl.State;
import rddl.competition.Client;
import rddl.competition.Server;
import rddl.parser.parser;
import rddl.policy.Policy;
import rddl.sim.Result;
import rddl.sim.Simulator;
import rddl.viz.StateViz;
import rddl.viz.TrafficDisplay;
import util.Pair;

public class TrafficSimulator {

	//RDDL stuff, in case we want to execute policy with RDDL server
	private TrafficDisplay _viz;
	
	private String _instName;

	private RDDL _rddl = new RDDL();
	
	private DOMAIN _d;
	private INSTANCE _i;
	private NONFLUENTS _n;
	
	private TrafficState _state;

	//local simulator, basically re-implements the RDDL simulator efficiently
	
	private Random _rand;
	
//	private Map< Integer, Map< TrafficState , TrafficState > > _noOpCache
//		= new HashMap<Integer, Map< TrafficState, TrafficState > >();
	
	private int _carSteps = 1;
	
	
	public TrafficSimulator( String instanceName, int nrows, int ncols, int blockSize, int wingSize, 
			int delayMillis, Policy pi, long rand_seed, int carSteps) throws Exception{
		
		_instName = instanceName;
		
		_viz = new TrafficDisplay(delayMillis, nrows, ncols, blockSize, wingSize);
		
		_rddl.addOtherRDDL( parser.parse(new File("traffic_mdp.rddl" )) );
		
		_rddl.addOtherRDDL( parser.parse(new File( _instName + ".rddl" ) ) );
		
		_d = _rddl._tmDomainNodes.values().iterator().next();
		
		_i = _rddl._tmInstanceNodes.values().iterator().next();
		
		_n = _rddl._tmNonFluentNodes.values().iterator().next();
		
		State rddlState = new State();
		
		rddlState.init(_n != null ? _n._hmObjects : null, _i._hmObjects,  
				_d._hmTypes, _d._hmPVariables, _d._hmCPF,
				_i._alInitState, _n == null ? null : _n._alNonFluents, 
						_d._alStateConstraints, _d._exprReward, _i._nNonDefActions);
		
		_state = new TrafficState(rddlState);

		_carSteps = carSteps;

		_rand = new Random( rand_seed );

		
	}
	
	//convenient method to set a state from RDDL
	//will be useful if we execute a policy against RDDL
	//not necessary for policy gradient
	public void setState( State rddlState ){
		_state = null;
		_state.setState(rddlState);
	}
	
	public Pair< TrafficState, Double > takeAction( ArrayList<PVAR_INST_DEF> aJointAction ){
		
		
		TrafficState oneStepNoOp = lookupNoOpCache( );

		if( oneStepNoOp == null ){
			//simulate noop
			oneStepNoOp = simulateNoOp( );
		}
		
		//first, change the light state
		//flip affects only light state

		State rddlState = _state.getState();
		
		for( PVAR_INST_DEF assign : aJointAction ){
			assert( rddlState.setPVariableAssign(assign._sPredName, assign._alTerms, assign._oValue) );
		}
		
		ArrayList<LCONST> temp = new ArrayList<LCONST>(1);
		
		HashMap<LVAR, LCONST> subs = new HashMap<LVAR,LCONST>();
		
		LVAR xionLvar = new LVAR("?i");
		
		for( LCONST xion : intersections ){
			//get cpt
			
			subs.clear();
		
			subs.put( xionLvar, xion);
			//substitute
			
			//update
			temp.set(0, xion);

			update( light1_primed, subs, temp );
			update( light2_primed, subs, temp );
			
		}
		
		//unset action
		for( PVAR_INST_DEF assign : aJointAction ){
			assert( rddlState.setPVariableAssign(assign._sPredName, assign._alTerms, null) );
		}
		
		int remaining = _carSteps-1;
		
		while( remaining > 0 ){
			
			for( int trying = remaining; trying > 0 && remaining > 0; --trying ){

				TrafficState lookup = lookupCache( trying );
				
				if( lookup != null ){
					remaining -= trying;
					trying = remaining + 1;
					cacheHit += trying;
				}
				
			}
			
			if( remaining > 0 ){
				simulateNoOp( );
				--remaining;
			}

		}
		//nothing was found
		//just simulate
		
		
		rddlState.computeNextState(aJointAction, _rand);
		
	}

	private void update( PVAR_NAME pvar, HashMap<LVAR,LCONST> subs, ArrayList<LCONST> temp ){

		
		try {
			State rddlState = _state.getState();
			CPF_DEF cpt = rddlState._hmCPFs.get(pvar);
			EXPR expr = cpt._exprEquals;
			boolean newVal;
			newVal = (Boolean) expr.sample(subs, rddlState, _rand);
			assert( rddlState.setPVariableAssign(pvar, temp, newVal) );
		} catch (EvalException e) {
			e.printStackTrace();
		}
		
		
	}
	
	public State getcurrentState(){
		return _state.getState();
	}
	
	public static void main(String[] args) throws Exception {
		TrafficSimulator r = new TrafficSimulator("traffic_mdp_3_3", 3, 3, 5, 2, 3000,  
				new AdvancePolicy(), 0 );
		
		Thread t = new Thread(r);
		t.start();
		t.join();
		
	}
	
}
